import os
import sys
from PIL import Image, ImageOps
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QPushButton, QFileDialog, QScrollArea, QMessageBox
)
from PySide6.QtGui import QPixmap, QImage
from PySide6.QtCore import Qt
import cv2  # 用于视频预览生成
import pillow_heif


def _generate_video_thumb(path, thumb_size):
    """生成视频缩略图（第一帧）"""
    try:
        cap = cv2.VideoCapture(path)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame)
            img.thumbnail(thumb_size)
            return QPixmap.fromImage(
                QImage(img.tobytes(), img.width, img.height, img.width * 3, QImage.Format.Format_RGB888)
            )
        cap.release()
    except Exception as e:
        print(f"Error loading video {path}: {e}")
        return None


def _generate_image_thumb(path, thumb_size):
    """生成图片缩略图"""
    try:
        img = Image.open(path)
        img = ImageOps.exif_transpose(img)  # 正确处理方向
        img.thumbnail(thumb_size)
        return QPixmap.fromImage(
            QImage(img.tobytes(), img.width, img.height, img.width * 3, QImage.Format.Format_RGB888)
        )
    except Exception as e:
        print(f"Error loading image {path}: {e}")
        return None


class LivePhotoPair:
    def __init__(self, image_path, video_path):
        self.image_path = image_path
        self.video_path = video_path
        self.image_thumb = None
        self.video_thumb = None

    def generate_thumbs(self, thumb_size=(256, 256)):
        """生成图片和视频的缩略图"""
        self.image_thumb = _generate_image_thumb(self.image_path, thumb_size)
        self.video_thumb = _generate_video_thumb(self.video_path, thumb_size)


class PhotoWidget(QWidget):
    def __init__(self, pair: LivePhotoPair, parent: 'MainWindow'):  # 明确指定 parent 类型
        super().__init__(parent)
        self.pair = pair
        self.parent_widget = parent  # 直接引用父组件
        self.init_ui()

    def init_ui(self):
        """初始化界面"""
        layout = QHBoxLayout()

        # 图片部分
        image_layout = self._create_media_layout("图片", self.pair.image_path, self.pair.image_thumb)
        layout.addLayout(image_layout)

        # 视频部分
        video_layout = self._create_media_layout("视频", self.pair.video_path, self.pair.video_thumb)
        layout.addLayout(video_layout)

        # 操作按钮
        btn_layout = QVBoxLayout()
        self.keep_both_btn = QPushButton("保留两者")
        self.keep_image_btn = QPushButton("仅保留图片")
        self.keep_video_btn = QPushButton("仅保留视频")

        btn_layout.addWidget(self.keep_both_btn)
        btn_layout.addWidget(self.keep_image_btn)
        btn_layout.addWidget(self.keep_video_btn)

        layout.addLayout(btn_layout)
        self.setLayout(layout)

        # 连接信号
        self.keep_image_btn.clicked.connect(lambda: self.parent_widget.process_pair(self.pair, "image"))
        self.keep_video_btn.clicked.connect(lambda: self.parent_widget.process_pair(self.pair, "video"))
        self.keep_both_btn.clicked.connect(lambda: self.parent_widget.process_pair(self.pair, "both"))

    def _create_media_layout(self, title, path, thumb):
        """创建图片或视频的布局"""
        layout = QVBoxLayout()
        label_name = QLabel(title)
        label_name.setAlignment(Qt.AlignCenter)
        layout.addWidget(label_name)

        label_media = QLabel()
        if thumb:
            label_media.setPixmap(thumb)
        layout.addWidget(label_media)

        label_path = QLabel(os.path.basename(path))
        label_path.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(label_path)

        return layout


class MainWindow(QMainWindow):
    def __init__(self, param_dir_path=None):
        super().__init__()
        self.dir_path = param_dir_path
        self.pairs = []
        self.init_ui()
        self.setWindowTitle("Live Photo 管理工具")
        self.resize(800, 600)

    def init_ui(self):
        """初始化主界面"""
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout()
        main_widget.setLayout(layout)

        # 顶部按钮
        top_layout = QHBoxLayout()
        self.btn_open = QPushButton("选择文件夹" if not self.dir_path else "读取" + self.dir_path)
        self.btn_open.clicked.connect(self.load_directory)
        top_layout.addWidget(self.btn_open)
        layout.addLayout(top_layout)

        # 滚动区域
        scroll = QScrollArea()
        self.content_widget = QWidget()
        self.scroll_layout = QVBoxLayout()
        self.content_widget.setLayout(self.scroll_layout)
        scroll.setWidget(self.content_widget)
        scroll.setWidgetResizable(True)
        layout.addWidget(scroll)

    def load_directory(self):
        """加载文件夹并扫描文件"""
        dir_path_temp = QFileDialog.getExistingDirectory(
            self, caption="选择照片目录") if not self.dir_path else self.dir_path
        if not dir_path_temp:
            return
        print(f"读取目录: {dir_path_temp}")
        self.btn_open.setEnabled(False)
        self.btn_open.setText(f"读取中: {dir_path_temp}")

        # 扫描文件
        files = os.listdir(dir_path_temp)
        base_names = {}
        for f in files:
            base, ext = os.path.splitext(f)
            ext = ext.lower()
            if ext in ['.jpg', '.jpeg', '.heic']:
                base_names.setdefault(base, {})['image'] = os.path.join(dir_path_temp, f)
            elif ext == '.mov':
                base_names.setdefault(base, {})['video'] = os.path.join(dir_path_temp, f)

        # 创建配对
        self.pairs = []
        for base, paths in base_names.items():
            if 'image' in paths and 'video' in paths:
                pair = LivePhotoPair(paths['image'], paths['video'])
                pair.generate_thumbs()
                self.pairs.append(pair)

        # 显示结果
        self.show_pairs()
        self.btn_open.setEnabled(True)
        self.btn_open.setText("选择文件夹")

    def show_pairs(self):
        """显示所有配对"""
        # 清空现有内容
        for i in reversed(range(self.scroll_layout.count())):
            self.scroll_layout.itemAt(i).widget().setParent(None)

        # 添加新内容
        for pair in self.pairs:
            widget = PhotoWidget(pair, self)  # 将 self 作为 parent 传递
            self.scroll_layout.addWidget(widget)

    def process_pair(self, pair, action):
        """处理文件操作"""
        if action == "image":
            os.remove(pair.video_path)
        elif action == "video":
            os.remove(pair.image_path)

        # 从列表中移除
        self.pairs.remove(pair)
        self.show_pairs()  # 刷新界面

        QMessageBox.information(self, "操作完成", "文件处理成功！")


def determine_dir()-> str | None:
    try:
        dir_path = sys.argv[1] if len(sys.argv) > 1 else None
        if dir_path and not os.path.isdir(dir_path):
            raise FileNotFoundError(f"目录不存在: {dir_path}")
        return dir_path
    except Exception as e:
        print(f"错误: {e}")
        return None


if __name__ == "__main__":
    app = QApplication(sys.argv)
    pillow_heif.register_heif_opener()  # 支持 HEIC 格式
    window = MainWindow(param_dir_path=determine_dir())
    window.show()
    sys.exit(app.exec())